import numpy as np
from traits.api import HasTraits, Instance, Button, \
    on_trait_change
from traitsui.api import View, Item, HSplit, Group
from mayavi import mlab
from mayavi.core.ui.api import MayaviScene, MlabSceneModel, SceneEditor
from traits.api import *
from traitsui.api import *
from mayavi.core.api import Engine

pi = np.pi
cos = np.cos
sin = np.sin
phi, theta = np.mgrid[0:pi:180j, 0:2 * pi:180j]
m0 = 4
m1 = 3
m2 = 2
m3 = 3
m4 = 1
m5 = 2
m6 = 2
m7 = 4
s = sin(m0 * phi) ** m1 + cos(m2 * phi) ** m3 + sin(m4 * theta) ** m5 + cos(m6 * theta) ** m7
x = sin(phi) * cos(theta)
y = cos(phi)
z = sin(phi) * sin(theta)

class Form(HasTraits):
    engine = Instance(Engine, ())
    scene1 = Instance(MlabSceneModel)
    scene2 = Instance(MlabSceneModel)
    button1 = Button('Redraw')
    button2 = Button('Redraw')
    def traits_init(self):
        self.engine.start()
    def _scene1_default(self):
        scene1 = MlabSceneModel(engine=self.engine)
        return scene1
    def _scene2_default(self):
        # self.engine.start()
        scene2 = MlabSceneModel(engine=self.engine)
        return scene2
    def picker_callback(self, picker_obj):
        picked = picker_obj.actors
        if self.mesh.actor.actor._vtk_obj in [o._vtk_obj for o in picked]:
           self.x_, self.y_ = np.lib.index_tricks.unravel_index(picker_obj.point_id,
                                                       s.shape)
           print("Data indices: %i, %i" % (self.x_, self.y_))
           self.n_x, self.n_y = s.shape
           self.cursor.mlab_source.reset(x=self.x_ - self.n_x / 2.,
                                     y=self.y_ - self.n_y / 2.)
           self.cursor3d.mlab_source.reset(x=x[self.x_, self.y_],
                                       y=y[self.x_, self.y_],
                                       z=z[self.x_, self.y_])
    @observe('scene1.activated')
    def redraw_scene1(self, event=None):
        self.mesh = self.scene1.mlab.mesh(x, y, z, scalars=s)
        self.cursor3d = self.scene1.mlab.points3d(0., 0., 0., mode='axes',
                                 color=(0, 0, 0),
                                 scale_factor=0.5)
        # self.scene1.mlab.title('Click on the ball')
        self.scene1.mlab.gcf().on_mouse_pick(self.picker_callback)

    @observe('scene2.activated')
    def redraw_scene2(self, event=None):
        im = self.scene2.mlab.imshow(s)
        self.cursor = self.scene2.mlab.points3d(0, 0, 0, mode='2dthick_cross',
                               color=(0, 0, 0),
                               scale_factor=10)
        self.scene2.mlab.view(60, 60)

    view = View(HSplit(
                  Group(
                       Item('scene1',
                            editor=SceneEditor(scene_class=MayaviScene), height=250,
                            width=300),
                       'button1',
                       show_labels=False,
                  ),
                  Group(
                       Item('scene2',
                            editor=SceneEditor(scene_class=MayaviScene), height=250,
                            width=300, show_label=False),
                       'button2',
                       show_labels=False,
                  ),
                ),
                resizable=True,
                )

from PyQt5.QtWidgets import *
class FormWidget(QWidget):
    def __init__(self):
        QWidget.__init__(self)
        layout = QHBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)
        self.form = Form()
        layout.addWidget(self.form.edit_traits(parent=self, kind='subpanel').control)

formWidget = FormWidget()
if __name__ == '__main__':
    import os
    import sys
    app = QApplication(sys.argv)
    formWidget.show()
    sys.exit(app.exec_())

